from typing import List
import os
import sys
import signal
import json
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import rospy
import actionlib
from tf.transformations import quaternion_from_euler
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
from geometry_msgs.msg import Pose, PoseStamped


def signal_handler(sig, frame):
    print("接收到信号，程序即将退出...")
    sys.exit(0)


signal.signal(signal.SIGINT, signal_handler)  # 注册 Ctrl+C 信号处理函数


def xy_yaw_to_pose(x, y, yaw):
    """
    Convert x, y, yaw to Pose
    """
    pose = Pose()
    pose.position.x = x
    pose.position.y = y
    q = quaternion_from_euler(0.0, 0.0, yaw)
    pose.orientation.x = q[0]
    pose.orientation.y = q[1]
    pose.orientation.z = q[2]
    pose.orientation.z = q[3]
    return pose


def convert_dict_to_pose(dict):
    # 将字典转换为Pose对象
    if "position" in dict and "orientation" in dict:
        if (
            "x" in dict["position"]
            and "y" in dict["position"]
            and "z" in dict["position"]
            and "w" in dict["orientation"]
        ):
            pose = Pose()
            pose.position.x = float(dict["position"]["x"])
            pose.position.y = float(dict["position"]["y"])
            pose.position.z = float(dict["position"]["z"])
            pose.orientation.w = float(dict["orientation"]["w"])
            pose.orientation.x = float(dict["orientation"]["x"])
            pose.orientation.y = float(dict["orientation"]["y"])
            pose.orientation.z = float(dict["orientation"]["z"])
        return pose
    return None


class SemanticMap:
    def __init__(self, waypoint_info_paths):
        with open(waypoint_info_paths, "r") as f:
            self.waypoints_info = json.load(f)

    def get_object_pose(self, object_name):
        """
        get the pose of the object
        """
        for scene, objects in self.waypoints_info.items():
            if object_name in objects:
                return objects[object_name]

        return None
    
    def load_waypoint_info(self):
        """
        load the waypoint information
        """
        return "{}".format(self.waypoints_info)


class Navigator:
    def __init__(self) -> None:
        # Create action client
        self.move_base_client = actionlib.SimpleActionClient('move_base', MoveBaseAction)
        self.pose_queue = []

    def navigate(self, poses: List[PoseStamped]):
        print('Waiting for move_base server...')
        self.move_base_client.wait_for_server()
        print('Connected to move_base server.')
        self.pose_queue = poses[1:]
        goal = MoveBaseGoal()
        goal.target_pose = poses[0]
        # Send the goal
        self.move_base_client.send_goal(goal, done_cb=self.done_cb, active_cb=self.active_cb, feedback_cb=self.feedback_cb)
        
    def cancel(self):
        self.pose_queue.clear()
        self.move_base_client.cancel_all_goals()
    
    def done_cb(self, status, result):
        if status == actionlib.GoalStatus.SUCCEEDED:
            print('Goal succeeded.')
            if len(self.pose_queue) > 0:
                next_goal = MoveBaseGoal()
                next_goal.target_pose = self.pose_queue[0]
                self.navigate(self.pose_queue)
            else:
                print('All goals reached.')
        else:
            print('Goal failed with status: %d' % status)

    def active_cb(self):
        # print('Goal is now active.')
        return
    
    def feedback_cb(self, feedback):
        # print('Received feedback: %s' % feedback)
        return


class LlmNavigatorInMap(Navigator):
    def __init__(self, map_path=None) -> None:
        super().__init__()
        self.semantic_map = SemanticMap(map_path)  # put the path of the waypoint information file here
        
        system_template = "You are a navigation system for robotics. You need to analysis which object waypoint should the robot go to according to the user's input and the waypoints information. The waypoints information are organized in json format, the first level is scene, and the second level is object waypoint. Only output the name of the object that best matches the user's input to me, no other things."
        user_template = "The information of the waypoints are: {waypoint_info}"
        prompt_template = ChatPromptTemplate.from_messages([
            ('system', system_template),
            ('user', user_template),
            ('user', '{task}')
        ])
        partial_template = prompt_template.partial(waypoint_info=self.semantic_map.load_waypoint_info())
        model = ChatOpenAI(model="gpt-4-1106-preview")
        self.chain = partial_template | model | StrOutputParser()
    
    def navigate_to_object(self, response_str):
        """
        find the corresponding pose
        use move base to navigate
        """
        pose_dict = self.semantic_map.get_object_pose(response_str)
        if pose_dict is not None:
            # 设置目标位置
            pose_stamped = PoseStamped()
            pose_stamped.header.frame_id = "map"
            pose_stamped.pose = convert_dict_to_pose(pose_dict)
            print("Going to: {}".format(response_str))
            print(pose_stamped)
            # 发送目标指令
            self.navigate([pose_stamped])
            return True
        
        print(f"{response_str}\nObject not found")
        return False

    def invoke(self, command):
        response = self.chain.invoke({"task": command})
        return self.navigate_to_object(response)


def main():
    # set up ros configuration
    rospy.init_node("llm_navigator", anonymous=True)
    llm_navigator = LlmNavigatorInMap(os.path.join(os.path.dirname(__file__), "../json/map.json"));
    while not rospy.is_shutdown():
        user_input = input("User: ")
        if user_input.lower() in ["exit", "quit", "bye"]:
            print("Chat session exit!")
            break
        llm_navigator.invoke(user_input)


if __name__ == "__main__":
    main()